package com.simple.core.config.mdc;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;

@Slf4j
public class MDCThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

    private static final String REQUEST_ID_KEY = MDCConstant.REQUEST_ID_KEY;

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        Map<String, String> context = MDC.getCopyOfContextMap();
        return super.submit(() -> {
            // 将父线程的MDC内容传给子线程
            T result = null;
            if (CollUtil.isNotEmpty(context) && StrUtil.isNotBlank(context.get(REQUEST_ID_KEY))) {
                MDC.setContextMap(context);
            } else {
                MDC.put(REQUEST_ID_KEY, IdUtil.fastSimpleUUID()); //为空设置新值
            }
            try {
                result = task.call();
            } catch (Exception e1) {
                log.error("mdc submit exception.", e1);
            } finally {
                try {
                    MDC.clear();
                } catch (Exception e2) {
                    log.warn("mdc clear exception.", e2);
                }
            }
            return result;
        });
    }

    @Override
    public void execute(Runnable task) {
        Map<String, String> context = MDC.getCopyOfContextMap();
        super.execute(() -> {
            if (CollUtil.isNotEmpty(context) && StrUtil.isNotBlank(context.get(REQUEST_ID_KEY))) {
                MDC.setContextMap(context);
            } else {
                MDC.put(REQUEST_ID_KEY, IdUtil.fastSimpleUUID()); //为空设置新值
            }
            try {
                task.run();
            } catch (Exception e1) {
                log.error("mdc submit exception.", e1);
            } finally {
                try {
                    MDC.clear();
                } catch (Exception e2) {
                    log.warn("mdc clear exception.", e2);
                }
            }
        });
    }
}
